from fastapi import APIRouter, HTTPException, Query, Request
from fastapi.responses import StreamingResponse, Response
import httpx
import io
from typing import Optional
import logging

router = APIRouter()
logger = logging.getLogger(__name__)

@router.get("/proxy")
@router.head("/proxy")
async def proxy_image(
    request: Request,
    url: str = Query(..., description="The image URL to proxy"),
    cache_control: Optional[str] = Query("public, max-age=3600", description="Cache control header")
):
    """
    代理外部图片请求，解决CORS问题
    
    Args:
        url: 要代理的图片URL
        cache_control: 缓存控制头
    
    Returns:
        StreamingResponse: 图片数据流
    """
    try:
        # 验证URL是否为允许的域名
        allowed_domains = [
            "trae-api-us.mchost.guru",
            "images.unsplash.com",
            "via.placeholder.com",
            "picsum.photos",
            "source.unsplash.com"
        ]
        
        # 检查URL是否来自允许的域名
        is_allowed = any(domain in url for domain in allowed_domains)
        if not is_allowed:
            raise HTTPException(status_code=400, detail="Domain not allowed")
        
        logger.info(f"Proxying image from: {url}")
        
        # 使用httpx获取图片
        async with httpx.AsyncClient(timeout=30.0, follow_redirects=True) as client:
            response = await client.get(url)
            
            if response.status_code != 200:
                logger.error(f"Failed to fetch image: {response.status_code}")
                raise HTTPException(
                    status_code=response.status_code, 
                    detail=f"Failed to fetch image: {response.status_code}"
                )
            
            # 获取内容类型
            content_type = response.headers.get("content-type", "image/jpeg")
            
            # 确保是图片类型
            if not content_type.startswith("image/"):
                logger.error(f"Invalid content type: {content_type}")
                raise HTTPException(status_code=400, detail="URL does not point to an image")
            
            # 准备响应头
            headers = {
                "Cache-Control": cache_control,
                "Access-Control-Allow-Origin": "*",
                "Access-Control-Allow-Methods": "GET, HEAD",
                "Access-Control-Allow-Headers": "*",
                "Content-Type": content_type
            }
            
            # 如果是HEAD请求，只返回头部信息
            if request.method == "HEAD":
                return Response(headers=headers)
            
            # 返回流式响应
            return StreamingResponse(
                io.BytesIO(response.content),
                media_type=content_type,
                headers=headers
            )
            
    except httpx.TimeoutException:
        logger.error(f"Timeout while fetching image: {url}")
        raise HTTPException(status_code=408, detail="Request timeout")
    except httpx.RequestError as e:
        logger.error(f"Request error while fetching image: {e}")
        raise HTTPException(status_code=502, detail="Failed to fetch image")
    except Exception as e:
        logger.error(f"Unexpected error: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")

@router.get("/trae-image")
@router.head("/trae-image")
async def proxy_trae_image(
    request: Request,
    prompt: str = Query(..., description="Image generation prompt"),
    image_size: str = Query("landscape_4_3", description="Image size")
):
    """
    专门代理Trae API图片生成请求
    
    Args:
        prompt: 图片生成提示词
        image_size: 图片尺寸
    
    Returns:
        StreamingResponse: 生成的图片数据流
    """
    try:
        # 构建Trae API URL
        trae_url = f"https://trae-api-us.mchost.guru/api/ide/v1/text_to_image?prompt={prompt}&image_size={image_size}"
        
        logger.info(f"Generating image via Trae API: {trae_url}")
        
        # 使用httpx获取图片
        async with httpx.AsyncClient(timeout=60.0, follow_redirects=True) as client:  # 增加超时时间，因为图片生成可能需要更长时间
            response = await client.get(trae_url)
            
            if response.status_code != 200:
                logger.error(f"Trae API failed: {response.status_code}")
                raise HTTPException(
                    status_code=response.status_code, 
                    detail=f"Trae API failed: {response.status_code}"
                )
            
            # 获取内容类型
            content_type = response.headers.get("content-type", "image/jpeg")
            
            # 准备响应头
            headers = {
                "Cache-Control": "public, max-age=86400",  # 缓存24小时
                "Access-Control-Allow-Origin": "*",
                "Access-Control-Allow-Methods": "GET, HEAD",
                "Access-Control-Allow-Headers": "*",
                "Content-Type": content_type
            }
            
            # 如果是HEAD请求，只返回头部信息
            if request.method == "HEAD":
                return Response(headers=headers)
            
            # 返回流式响应
            return StreamingResponse(
                io.BytesIO(response.content),
                media_type=content_type,
                headers=headers
            )
            
    except httpx.TimeoutException:
        logger.error(f"Timeout while generating image via Trae API")
        raise HTTPException(status_code=408, detail="Image generation timeout")
    except httpx.RequestError as e:
        logger.error(f"Request error while calling Trae API: {e}")
        raise HTTPException(status_code=502, detail="Failed to generate image")
    except Exception as e:
        logger.error(f"Unexpected error in Trae API proxy: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")